'''
    对网络的每层输出和卷积核参数进行可视化
'''
import os
import numpy as np
import cv2
from PIL import Image
import torch
from torchsummary import summary
from nets import get_model_from_name
from utils.utils import get_classes, predict_input, makedir



class NetVisualize():
    _defaults = {
        # --------------------------------------------------------------------------#
        #   使用自己训练好的模型进行预测一定要修改model_path和classes_path！
        #   model_path指向logs文件夹下的权值文件，classes_path指向model_data下的txt
        #   如果出现shape不匹配，同时要注意训练时的model_path和classes_path参数的修改
        # --------------------------------------------------------------------------#
        "model_path": 'logs/best_epoch_weights.pth',
        # "model_path": 'logs/ep100-loss0.000-val_loss0.000.pkl',
        "classes_path": 'model_data/news_classes.txt',
        # --------------------------------------------------------------------#
        #   输入的图片大小
        # --------------------------------------------------------------------#
        "input_shape": [224, 224],
        # --------------------------------------------------------------------#
        #   所用模型种类：
        #   LeNet、AlexNet、vgg16、mobilenetv2、resnet50、
        #   mobilenetv3
        # --------------------------------------------------------------------#
        "backbone": 'mobilenetv3',
        # -------------------------------#
        #   是否使用Cuda
        #   没有GPU可以设置成False
        # -------------------------------#
        "cuda": True,
        # --------------------------------------#
        #   要保存的特征层和卷积层可视化图片的文件夹
        # --------------------------------------#
        "save_path": 'net_visualize'
    }

    # ---------------------------------------------------#
    #   初始化classification
    # ---------------------------------------------------#
    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)  # 把_defaults中的参数，初始化为类的属性
        for name, value in kwargs.items():
            setattr(self, name, value)  # 设置类的属性

        # ---------------------------------------------------#
        #   获得种类
        # ---------------------------------------------------#
        self.class_names, self.num_classes = get_classes(self.classes_path)
        self.generate()

    # ---------------------------------------------------#
    #   获得所有的分类
    # ---------------------------------------------------#
    def generate(self):
        # ---------------------------------------------------#
        #   载入模型与权值
        # ---------------------------------------------------#
        if self.backbone == "LeNet":
            self.model = get_model_from_name[self.backbone](num_classes=self.num_classes, in_channels=3)
        elif self.backbone == "mobilenetv3":
            self.model = get_model_from_name[self.backbone](n_class=self.num_classes, mode='large')
        else:
            self.model = get_model_from_name[self.backbone](num_classes=self.num_classes)
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model.load_state_dict(torch.load(self.model_path, map_location=device))
        self.model = self.model.eval()
        print('{} model, and classes loaded.'.format(self.model_path))

        if self.cuda:
            self.model = self.model.cuda()

    # ---------------------------------------------------#
    #   获得模型信息及参数大小
    # ---------------------------------------------------#
    def get_model_info(self):
        print("get_model_info:")
        device = 'cuda' if self.cuda else 'cpu'
        summary(model=self.model, input_size=(3, self.input_shape[0], self.input_shape[1]), device=device)

    # --------------------------------------------#
    #   输入Image类型的数据，提取每一层网络的输出特征层
    #   therd_size指定保存图像的大小
    # --------------------------------------------#
    def get_feature_output(self, image, therd_size):
        # ---------------------------------------------#
        #   对测试图片预处理
        # ---------------------------------------------#
        image = predict_input(image, self.input_shape)
        # ---------------------------------------------------------#
        #   添加batch_size维度
        # ---------------------------------------------------------#
        image = np.expand_dims(image, 0)

        with torch.no_grad():
            x = torch.from_numpy(image)
            if self.cuda:
                x = x.cuda()

        output_feature_layer = []

        # model.named_modules() 返回generator, 能够按顺序迭代网络的每一层及其子层以及名字
        for name, module in self.model.named_modules():
            # 如果该层不为嵌套层,进行前向传播
            if (
                    not isinstance(module, torch.nn.Sequential)
                    and not isinstance(module, torch.nn.ModuleList)
                    and not (module == self.model)
            ):
                print(name)
                print(module)

        # model._modules.items()获取网络的最外层及其名字
        # for name, module in self.model._modules.items():
        #     if "features" in name:
        #         for layer in module:
        #             with torch.no_grad():
        #                 x = layer(x)
        #             output_feature_layer.append(x)

        # 对获得的特征层列表进行遍历
        for i, imagedata in enumerate(output_feature_layer):
            features = imagedata[0]  # 去除batch_size维度
            channels = features.shape[0]  # 得到该特征层的通道数

            dst_path = os.path.join(self.save_path, "features", "layer_{}".format(i))
            makedir(dst_path)

            # 对特征层的每个通道遍历
            for c in range(channels):
                feature = features.data.cpu().numpy()
                feature_img = feature[c, :, :]
                feature_img = np.asarray(feature_img * 255, dtype=np.uint8)

                # 对图像进行伪彩色增强,热力图
                feature_img = cv2.applyColorMap(feature_img, cv2.COLORMAP_JET)

                # 如果图片尺寸小于therd_size, 将图像resize到therd_size尺寸
                if feature_img.shape[0] < therd_size:
                    tmp_file = os.path.join(dst_path, "channel_{}_{}.jpg".format(c, therd_size))
                    tmp_img = feature_img.copy()
                    tmp_img = cv2.resize(tmp_img, (therd_size, therd_size), interpolation=cv2.INTER_NEAREST)
                    cv2.imwrite(tmp_file, tmp_img)
                else:
                    dst_file = os.path.join(dst_path, "channel_{}.jpg".format(c))
                    cv2.imwrite(dst_file, feature_img)

    # --------------------------------------------#
    #   对模型的卷积核进行可视化
    #   therd_size指定保存图像的大小
    # --------------------------------------------#
    def get_conv_visual(self, therd_size):
        # 遍历模型中每一层网络的参数
        for name, parameters in self.model.named_parameters():
            # 如果参数的维度为4,说明是卷积层
            if parameters.ndim == 4:

                dst_path = os.path.join(self.save_path, "conv_visual", "layer_{}".format(name))
                makedir(dst_path)

                # 卷积核的个数
                kernel_number = parameters.shape[0]
                # 卷积核的通道数
                kernel_channel = parameters.shape[1]
                # 遍历卷积核的个数
                for n in range(kernel_number):
                    kernel = parameters[n, :, :, :]
                    # 遍历卷积核的通道
                    for c in range(kernel_channel):
                        kernel_image = kernel[c, :, :]

                        feature_img = np.asarray(kernel_image.cpu().detach().numpy() * 255, dtype=np.uint8)
                        # 对图像进行伪彩色增强,热力图
                        feature_img = cv2.applyColorMap(feature_img, cv2.COLORMAP_JET)

                        # 如果图片尺寸小于therd_size, 将图像resize到therd_size尺寸
                        if feature_img.shape[0] < therd_size:
                            tmp_file = os.path.join(dst_path, "channel_{}_{}.jpg".format(c, therd_size))
                            tmp_img = feature_img.copy()
                            tmp_img = cv2.resize(tmp_img, (therd_size, therd_size), interpolation=cv2.INTER_CUBIC)
                            cv2.imwrite(tmp_file, tmp_img)
                        else:
                            dst_file = os.path.join(dst_path, "channel_{}.jpg".format(c))
                            cv2.imwrite(dst_file, feature_img)


if __name__ == "__main__":

    netVisualize = NetVisualize()
    netVisualize.get_model_info()
    # image = Image.open("datasets/test/cat/cat.10003.jpg")
    # netVisualize.get_feature_output(image, 256)
    # netVisualize.get_conv_visual(256)
